package step

import (
	"bytes"
	"github.com/wanchain/go-mpc/common"
	"github.com/wanchain/go-mpc/common/hexutil"
	"github.com/wanchain/go-mpc/log"
	"github.com/wanchain/go-mpc/p2p/discover"
	"github.com/wanchain/go-mpc/storeman/osmconf"
	mpcprotocol "github.com/wanchain/go-mpc/storeman/storemanmpc/protocol"
	"math/big"
	"strconv"
)

type MpcAckRSStepUni340 struct {
	BaseStep
	message    map[discover.NodeID]bool
	remoteMpcR map[discover.NodeID][]mpcprotocol.CurvePointer // R
	remoteMpcS map[discover.NodeID][]big.Int                  // S
	mpcR       []mpcprotocol.CurvePointer
	mpcS       []big.Int
	sigNum     int
}

func CreateAckMpcRSStepUni340(totalNum int, sigNum int) *MpcAckRSStepUni340 {
	mpc := &MpcAckRSStepUni340{
		*CreateBaseStep(totalNum, -1),
		make(map[discover.NodeID]bool),
		make(map[discover.NodeID][]mpcprotocol.CurvePointer),
		make(map[discover.NodeID][]big.Int),
		make([]mpcprotocol.CurvePointer, 0),
		make([]big.Int, 0),
		sigNum}

	return mpc
}

func (mars *MpcAckRSStepUni340) InitStep(result mpcprotocol.MpcResultInterface) error {
	log.SyslogDebug("MpcAckRSStepUni340.InitStep begin", "ctx", mars.GetMpcContextId())
	mars.BaseStep.InitStep(result)

	for i := 0; i < mars.sigNum; i++ {
		key := mpcprotocol.RPk + strconv.Itoa(i)
		rpkBytes, err := result.GetByteValue(key)
		pt, err := mars.schnorrMpcer.UnMarshPt(rpkBytes)
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::InitStep", "ack mpc account step, init fail UnMarshPt. err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}
		mars.mpcR = append(mars.mpcR, pt)

		key = mpcprotocol.MpcS + strconv.Itoa(i)
		sValue, err := result.GetValue(key)
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::InitStep", "ack mpc account step, init fail. err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}
		mars.mpcS = append(mars.mpcS, sValue[0])
	}
	return nil
}

func (mars *MpcAckRSStepUni340) CreateMessage() []mpcprotocol.StepMessage {
	_, grpIdString, _ := mars.BaseStep.GetGrpId()
	isLeader, _ := osmconf.GetOsmConf().IsLeader(grpIdString)
	log.SyslogDebug("MpcAckRSStepUni340 CreateMessage", "leader ", isLeader, "ctx", mars.GetMpcContextId())
	if !isLeader {
		// not-leader no need to wait others message, only wait message from itself.
		mars.SetWaiting(1)
	}

	msg := mpcprotocol.StepMessage{
		MsgCode:   mpcprotocol.MPCMessage,
		PeerID:    nil,
		Data:      nil,
		BytesData: nil}

	msg.BytesData = make([][]byte, mars.sigNum)
	msg.Data = make([]big.Int, mars.sigNum)

	for i := 0; i < mars.sigNum; i++ {
		rpkBytes, err := mars.schnorrMpcer.MarshPt(mars.mpcR[i])
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::CreateMessage", "MarshPt err", err.Error(), "ctx", mars.GetMpcContextId(), "sigIndex", i)
		}
		log.SyslogInfo("MpcAckRSStepUni340::CreateMessage", "MarshPt result", hexutil.Encode(rpkBytes), "sigIndex", i, "ctx", mars.GetMpcContextId())

		msg.BytesData[i] = rpkBytes
		msg.Data[i] = mars.mpcS[i]
	}

	return []mpcprotocol.StepMessage{msg}
}

func (mars *MpcAckRSStepUni340) FinishStep(result mpcprotocol.MpcResultInterface, mpc mpcprotocol.StoremanManager) error {
	log.SyslogDebug("MpcAckRSStepUni340.FinishStep begin", "ctx", mars.GetMpcContextId())
	err := mars.BaseStep.FinishStep()
	if err != nil {
		return err
	}

	err = mars.verifyRS(result)
	if err != nil {
		return err
	}

	for i := 0; i < mars.sigNum; i++ {
		// rpk : R(x)
		var buffer bytes.Buffer
		xbytes, err := mars.schnorrMpcer.XBytes(mars.mpcR[i])
		if err != nil {
			return err
		}
		buffer.Write(xbytes)
		// S
		sFullBytes := common.LeftPadBytes(mars.mpcS[i].Bytes(), 32)
		//buffer.Write(mars.mpcS[i].Bytes())
		buffer.Write(sFullBytes[:])
		key := mpcprotocol.MpcContextResult + strconv.Itoa(i)
		result.SetByteValue(key, buffer.Bytes())
	}

	return nil
}

func (mars *MpcAckRSStepUni340) HandleMessage(msg *mpcprotocol.StepMessage) bool {
	log.SyslogInfo("MpcAckRSStepUni340.HandleMessage begin", "ctx", mars.GetMpcContextId())
	_, exist := mars.message[*msg.PeerID]
	if exist {
		log.SyslogErr("MpcAckRSStepUni340::HandleMessage", "MpcAckRSStep.HandleMessage fail. peer doesn't exist in task peer group. peerID",
			msg.PeerID.SlimString(), "ctx", mars.GetMpcContextId())
		return false
	}

	if len(msg.Data) == mars.sigNum && len(msg.BytesData) == mars.sigNum {

		for i := 0; i < mars.sigNum; i++ {
			mars.remoteMpcS[*msg.PeerID] = append(mars.remoteMpcS[*msg.PeerID], msg.Data[i])
			rpt, err := mars.schnorrMpcer.UnMarshPt(msg.BytesData[i])
			if err != nil {
				log.SyslogErr("MpcAckRSStepUni340::HandleMessage", "UnMarshPt error", err.Error(), "ctx", mars.GetMpcContextId())
				return false
			}
			mars.remoteMpcR[*msg.PeerID] = append(mars.remoteMpcR[*msg.PeerID], rpt)
		}

	}
	mars.message[*msg.PeerID] = true
	return true
}

func (mars *MpcAckRSStepUni340) verifyRS(result mpcprotocol.MpcResultInterface) error {
	// check R
	for _, mpcR := range mars.remoteMpcR {
		if mpcR == nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS remote R is nil", "ctx", mars.GetMpcContextId())
			return mpcprotocol.ErrInvalidMPCR
		}

		for i := 0; i < mars.sigNum; i++ {
			if !mars.schnorrMpcer.Equal(mars.mpcR[i], mpcR[i]) {
				localR := mars.mpcR[i]
				remoteR := mpcR[i]

				localRStr := mars.schnorrMpcer.PtToHexString(localR)
				remoteRStr := mars.schnorrMpcer.PtToHexString(remoteR)
				log.SyslogErr("MpcAckRSStepUni340::verifyRS remote R is NOT equal self R", "self R", localRStr, "remote R", remoteRStr, "ctx", mars.GetMpcContextId())

				return mpcprotocol.ErrInvalidMPCR
			}
		}

	}
	// check S
	for _, mpcS := range mars.remoteMpcS {

		for i := 0; i < mars.sigNum; i++ {
			if mars.mpcS[i].Cmp(&mpcS[i]) != 0 {
				localSStr := hexutil.Encode(mars.mpcS[i].Bytes())
				remoteSStr := hexutil.Encode(mpcS[i].Bytes())

				log.SyslogErr("MpcAckRSStepUni340::verifyRS remote S is NOT equal self S", "self S", localSStr, "remote S", remoteSStr, "ctx", mars.GetMpcContextId())
				return mpcprotocol.ErrInvalidMPCS
			}
		}

	}

	// check signVerify
	for i := 0; i < mars.sigNum; i++ {
		hashMBytes := mars.MpcContext.GetInput().HashData[i]

		gpkItem := mars.GetMpcContext().GetInput().PKBytes
		smpcer := mars.schnorrMpcer
		gpk, err := smpcer.UnMarshPt(gpkItem[:])
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "UnMarshPt err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}

		// rpk : R
		rpkBytes, err := smpcer.MarshPt(mars.mpcR[i])
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "MarshPt err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}

		xBytesRpk, err := smpcer.XBytesByBytes(rpkBytes)
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "XBytesByBytes xBytesRpk err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}

		XBytesGpk, err := smpcer.XBytesByBytes(gpkItem)
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "XBytesByBytes XBytesGpk  err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}
		m, err := mars.getm(i)
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "GetSchnorr340E err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}
		nSubE, err := smpcer.NSubX(m)
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "NSubX err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}

		// ssG + (n-e)*gpkLift = rpk
		ssG, err := smpcer.SkG(&mars.mpcS[i])
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "SkG err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}

		gpkLift, err := smpcer.LiftX(big.NewInt(0).SetBytes(XBytesGpk))
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "LiftX err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}

		//gpkLift*(n-e)
		temp, err := smpcer.MulPK(nSubE, gpkLift)
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "MulPK err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}

		// left
		left, err := smpcer.Add(ssG, temp)
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "Add err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}

		if !smpcer.IsOnCurve(left) {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS left is Not on curve", "ctx", mars.GetMpcContextId())
			return mpcprotocol.ErrNotOnCurve
		}

		xBytesLeft, err := smpcer.XBytes(left)
		if err != nil {
			log.SyslogErr("MpcAckRSStepUni340::verifyRS", "XBytes xBytesLeft err", err.Error(), "ctx", mars.GetMpcContextId())
			return err
		}

		log.Info("@@@@@@@@@@@@@@verifyRS 340@@@@@@@@@@@@@@",
			"m", hexutil.Encode(m.Bytes()),
			"R", hexutil.Encode(rpkBytes),
			"HashData", hexutil.Encode(hashMBytes[:]),
			"gpk(x)", hexutil.Encode(XBytesGpk),
			"R(x)", hexutil.Encode(xBytesRpk),
			"s", hexutil.Encode(mars.mpcS[i].Bytes()),
			"left (ssG + (n-e)*gpkLift)", smpcer.PtToHexString(left),
			"sG", smpcer.PtToHexString(ssG),
			"gpk", smpcer.PtToHexString(gpk),
			"ctx", mars.GetMpcContextId())

		if bytes.Equal(xBytesLeft, xBytesRpk) {
			log.SyslogInfo("MpcAckRSStepUni340::Verification success", "ctx", mars.GetMpcContextId())
		} else {
			log.SyslogErr("MpcAckRSStepUni340::Verification failed", "ctx", mars.GetMpcContextId())
			return mpcprotocol.ErrVerifyFailed
		}

	}

	return nil
}
